import argparse
import json
import os
import socket
import logging
import struct
import hashlib
from os.path import getsize, join

from banner.banner import MyBanner

OP_SAVE, OP_DELETE, OP_GET, OP_UPLOAD, OP_DOWNLOAD, OP_BYE, OP_LOGIN, OP_ERROR = 'SAVE', 'DELETE', 'GET', 'UPLOAD', 'DOWNLOAD', 'BYE', 'LOGIN', "ERROR"
TYPE_FILE, TYPE_DATA, TYPE_AUTH, DIR_EARTH = 'FILE', 'DATA', 'AUTH', 'EARTH'
FIELD_OPERATION, FIELD_DIRECTION, FIELD_TYPE, FIELD_USERNAME, FIELD_PASSWORD, FIELD_TOKEN = 'operation', 'direction', 'type', 'username', 'password', 'token'
FIELD_KEY, FIELD_SIZE, FIELD_TOTAL_BLOCK, FIELD_MD5, FIELD_BLOCK_SIZE = 'key', 'size', 'total_block', 'md5', 'block_size'
FIELD_STATUS, FIELD_STATUS_MSG, FIELD_BLOCK_INDEX = 'status', 'status_msg', 'block_index'
DIR_REQUEST, DIR_RESPONSE = 'REQUEST', 'RESPONSE'

logger = logging.getLogger('')


def set_logger(logger_name):
    """
    Create a logger
    :param logger_name: 日志名称
    :return: logger
    """
    logger_ = logging.getLogger(logger_name)  # 不加名称设置root logger
    logger_.setLevel(logging.INFO)

    formatter = logging.Formatter(
        '\033[0;34m%s\033[0m' % '%(asctime)s-%(name)s[%(levelname)s] %(message)s @ %(filename)s[%(lineno)d]',
        datefmt='%Y-%m-%d %H:%M:%S')

    # --> SCREEN DISPLAY
    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)
    ch.setFormatter(formatter)

    logger_.propagate = False
    logger_.addHandler(ch)
    return logger_


logger = set_logger('CLIENT')


class SocketClient:
    def __init__(self, ip, port):
        self.user_name = None
        self.file_path = None
        self.block_size = None
        self.total_block = None
        self.key = None
        self.token = None
        self.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.client.connect((ip, port))

    def close(self):
        logger.info(f'----关闭连接')
        self.client.close()

    def send(self, head, j_data, b_data=b''):
        b_str = head + j_data.encode() + b_data
        logger.info(f'----发送数据{b_str}')
        self.client.send(b_str)

    def recv(self):
        head = self.client.recv(8)
        j_len, b_len = struct.unpack('!II', head)
        j_data = json.loads(self.client.recv(j_len))
        logger.info(f'----接收数据j_data{j_data}')
        b_data = self.client.recv(b_len)
        logger.info(f'----接收数据b_data{b_data}')
        return j_data, b_data

    def login(self, user_name):
        logger.info(f'>---开始初始化登录.')
        self.user_name = user_name
        md5_password = hashlib.md5(self.user_name.encode()).hexdigest().lower()
        j_data = json.dumps({
            FIELD_OPERATION: OP_LOGIN,
            FIELD_DIRECTION: DIR_REQUEST,
            FIELD_TYPE: TYPE_AUTH,
            FIELD_USERNAME: self.user_name,
            FIELD_PASSWORD: md5_password
        })
        head = struct.pack('!II', len(j_data), 0)
        self.send(head, j_data)
        j_data, b_data = self.recv()
        self.token = j_data['token']
        logger.info(f'<---登录结束.')

    def save(self, file_path, key):
        logger.info(f'>---开始初始化save.')
        self.file_path = file_path
        file_size = getsize(file_path)
        j_data = json.dumps({
            FIELD_OPERATION: OP_SAVE,
            FIELD_DIRECTION: DIR_REQUEST,
            FIELD_TYPE: TYPE_FILE,
            FIELD_TOKEN: self.token,
            FIELD_SIZE: file_size,
            FIELD_KEY: key
        })
        head = struct.pack('!II', len(j_data), 0)
        self.send(head, j_data)
        j_data, b_data = self.recv()
        self.key = j_data[FIELD_KEY]
        self.total_block = j_data['total_block']
        self.block_size = j_data['block_size']
        logger.info(f'<---save结束.')

    def upload(self):
        logger.info(f'>---开始初始化upload.')
        for num in range(0, self.total_block):
            j_data = json.dumps({
                FIELD_OPERATION: OP_UPLOAD,
                FIELD_DIRECTION: DIR_REQUEST,
                FIELD_TYPE: TYPE_FILE,
                FIELD_TOKEN: self.token,
                FIELD_KEY: self.key,
                FIELD_BLOCK_INDEX: num
            })
            fp = open(self.file_path, 'rb+')
            fp.seek(self.block_size * num)
            b_data = fp.read(self.block_size)
            head = struct.pack('!II', len(j_data), len(b_data))
            self.send(head, j_data, b_data)
            j_data, b_data = self.recv()
        logger.info(f'<---upload结束.')

    def get(self, key):
        logger.info(f'>---开始初始化get.')
        os.makedirs(join('download', self.user_name), exist_ok=True)
        j_data = json.dumps({
            FIELD_OPERATION: OP_GET,
            FIELD_DIRECTION: DIR_REQUEST,
            FIELD_TYPE: TYPE_FILE,
            FIELD_TOKEN: self.token,
            FIELD_KEY: key
        })
        head = struct.pack('!II', len(j_data), 0)
        self.send(head, j_data)
        j_data, b_data = self.recv()
        self.key = j_data[FIELD_KEY]
        self.total_block = j_data['total_block']
        self.block_size = j_data['block_size']
        file_path = join('download', self.user_name, self.key)
        with open(file_path, 'wb+') as fid:
            fid.seek(j_data['size'] - 1)
            fid.write(b'\0')
        fid.close()
        logger.info(f'<---get结束.')

    def download(self):
        logger.info(f'>---开始初始化download.')
        file_path = join('download', self.user_name, self.key)
        for num in range(0, self.total_block):
            j_data = json.dumps({
                FIELD_OPERATION: OP_DOWNLOAD,
                FIELD_DIRECTION: DIR_REQUEST,
                FIELD_TYPE: TYPE_FILE,
                FIELD_TOKEN: self.token,
                FIELD_KEY: self.key,
                FIELD_BLOCK_INDEX: num
            })
            head = struct.pack('!II', len(j_data), 0)
            self.send(head, j_data)
            j_data, b_data = self.recv()
            with open(file_path, 'rb+') as fid:
                fid.seek(self.block_size * num)
                fid.write(b_data)
            fid.close()
        logger.info(f'<---download结束.')

    def delete(self, key):
        logger.info(f'>---开始初始化delete.')
        j_data = json.dumps({
            FIELD_OPERATION: OP_DELETE,
            FIELD_DIRECTION: DIR_REQUEST,
            FIELD_TYPE: TYPE_FILE,
            FIELD_TOKEN: self.token,
            FIELD_KEY: key
        })
        head = struct.pack('!II', len(j_data), 0)
        self.send(head, j_data)
        j_data, b_data = self.recv()
        logger.info(f'<---delete结束.')


def _argparse():
    parse = argparse.ArgumentParser()
    parse.add_argument("--user_name", default='wubo', action='store', required=False, dest="user_name",
                       help="用户名，默认wubo")
    parse.add_argument("--type", default='file', action='store', required=False, dest="type",
                       help="操作类型，默认文件类型")
    parse.add_argument("--operation", default='upload', action='store', required=False, dest="operation",
                       help="操作方式，默认上传文件")
    parse.add_argument("--file_path", default='example\\1.jpg', action='store', required=False, dest="file_path",
                       help="上传的文件地址/下载时保存文件地址")
    parse.add_argument("--key", default='1.jpg', action='store', required=False, dest="key",
                       help="文件唯一标识")
    return parse.parse_args()


def main():
    parser = _argparse()

    client = SocketClient('127.0.0.1', 1379)  # 初始化连接
    client.login(parser.user_name)  # 这里是用户名，随便，会按照用户名生成对应的文件夹进行文件存储

    if parser.type == 'file' and parser.operation == 'upload':
        client.save(parser.file_path, parser.key)  # 文件位置，上传后存储的key
        client.upload()  # 上传文件
        client.close()  # 关闭连接

    if parser.type == 'file' and parser.operation == 'download':
        os.makedirs('download', exist_ok=True)
        client.get(parser.key)  # 文件唯一标识
        client.download()  # 下载文件
        client.close()  # 关闭连接

    if parser.type == 'file' and parser.operation == 'delete':
        client.delete(parser.key)  # 文件唯一标识
        client.close()  # 关闭连接


if __name__ == '__main__':
    main()
